# A toolkit for microlensing calculations in R
# Pierre Vermaak, November 2008
# PLEASE NOTE I use a non-standard coordinate system for binary lenses:
# 1) The origin is always at the primary
# 2) Einstein radius always in units of the Primary, not the combined mass.
# 3) Full details in my thesis, available at: http://sites.google.com/site/examplebasedmicrolensing/

    lensTp<-function(t,te,tm) {
        return((t-tm)/te)
    }

    lensU<-function(t,te,tm,b) {
        return(sqrt(b^2+lensTp(t,te,tm)^2))
    }

    lensBlend<-function(f,a) {
        return(f * (a - 1.0) + 1.0)
    }

lensSLAmp<-function(t,b,te,tm,f) {
    u<-lensU(t,te,tm,b);
    amp<- (u * u + 2.0) / u / sqrt(u * u + 4.0);
    return(lensBlend(f,amp))
}

    lensMag<-function(a,m0){
        return(-2.5*log10(a)+m0)
    }

    lensMagFunc<-function(par){
        return(lensMag(lensSLAmp(par[1],par[2],par[5],par[6],par[3]),par[4]))
    }

    myChiSq<-function(x1,x2,err){
        return(sum(((x1-x2)/err)^2))
    }

lensChiSq<-function(t,y,err,b,f,m0,te,tm) {
    ymodel<-lensMag(lensSLAmp(t,b,te,tm,f),m0);
    return(myChiSq(y,ymodel,err))
}

    lensFit<-function(p) {
        return(lensChiSq(xdata,ydata,errdata,p[1],p[2],p[3],p[4],p[5]))
    }

rational1<-function(p,t) {
    return (1/(p[1]+p[2]*t+p[3]*t*t)**2);
}

rational1Sla <- function(p,t) {
    return(lensMag(lensSLAmp(t,p[1],p[4],p[5],p[2]),p[3]) + rational1(p[6:8],t));
}

rational1LensChiSq<-function(t,y,err,b,f,m0,te,tm,p) {
    ymodel<- lensMag(lensSLAmp(t,b,te,tm,f),m0) + rational1(p,t);
    return(myChiSq(y,ymodel,err))
}

rational1Fit<-function(p) {
    return(rational1LensChiSq(xdata,ydata,errdata,p[1],p[2],p[3],p[4],p[5],p[6:8]));
}

lensSourcePlanePosition <- function(day,a,ang,b,q,te,tm,rho=0,psi=0) {
# some constants
    lambda <- 4.726;
    tp <- 1099.0;
    beta <- -0.08;
    e <- 0.0167;

    aries <- 77. * pi / 180;
    funnyphi <- aries + lambda + pi;
    T <- 365.26;

# calculate x1tmax and x2tmax
    tptime <- (2 * pi * (tm - tp) / T);
    sqgl  <-  tptime + 2 * e * sin(tptime) - funnyphi;

    x1tmax  <-  rho * (-sin(beta) * cos(funnyphi) *
            (cos(sqgl) - e) - sin(beta) * sin(funnyphi) * sqrt(1. - e * e) * sin(sqgl));
    x2tmax  <-  rho * (-sin(funnyphi) * (cos(sqgl) - e) + cos(funnyphi) *
            (1. - e * e) * sin(sqgl));

# Now for today's point
    tptime  <-  (2 * pi * (day - tp) / T);
    p0  <-  (day - tm) / te;
    sqgl  <-  tptime + 2 * e * sin(tptime) - funnyphi;

    x1  <-  rho * (-sin(beta) * cos(funnyphi) *
            (cos(sqgl) - e) - sin(beta) * sin(funnyphi) * sqrt(1. - e * e) * sin(sqgl));
    x2  <-  rho * (-sin(funnyphi) * (cos(sqgl) - e) + cos(funnyphi) *
            (1. - e * e) * sin(sqgl));

    poft  <-  p0 + cos(psi) * (x1 - x1tmax) + sin(psi) * (x2 - x2tmax);
    doft  <-  b - sin(psi) * (x1 - x1tmax) + cos(psi) * (x2 - x2tmax);

# with prlx done, we convert poft and doft to the real values in the plane
    radang  <-  ang * pi / 180.;
    cosradangminusninety  <-  cos(radang - pi / 2.0);
    sinradangminusninety  <-  sin(radang - pi / 2.0);

    sx  <-  poft * cosradangminusninety - doft * sinradangminusninety;
    sy  <-  poft * sinradangminusninety + doft * cosradangminusninety;

    return (list("sx"=sx,"sy"=sy))
}

calcAsadaParams <- function(sx,sy,a,q) {
# scaling factor for all asada lengths
# Set M1 = 1, M2 = q
# Also, sx and sy will have to change
    scale = sqrt(1.0+q);

    nu = q/(1.0+q);
    l = a/scale;
    rho = sqrt(sx*sx + sy*sy)/scale;
    c = sx/rho/scale;
    s = sy/rho/scale;

    return(list("nu"=nu,"l"=l,"rho"=rho,"c"=c,"s"=s))
}


calcImagePositionsUsingAsada <- function(asadaPos, qForScaling)  {
    r = asadaPos$rho
        l = asadaPos$l
        c = asadaPos$c
        s = asadaPos$s
        v = asadaPos$nu

        r2 = r * r
        r3 = r2 * r
        r4 = r3 * r

        c2 = c * c
        c3 = c2 * c
        c4 = c3 * c

        s2 = s * s
        s3 = s2 * s
        s4 = s3 * s

        l2 = l * l
        l3 = l2 * l
        l4 = l3 * l

        a <- rep(0,6)

#/ start the coefficients proper
        a[6] = v * l * r3 * s3;
    a[5] = r2 * s2 + 2.0 * l * r3 * c * s2 - l2 * (2.0 * r2 * s2 - r4 * s2) - 2.0 * l3 * r3 * c * s2 + l4 * r2 * s2 - v * (5.0 * l * r3 * c * s2 - 4.0 * l2 * r2 * s2);
    a[4] = -2.0 * r2 * c * s - 4.0 * l * r3 * (c2 * s - s3) + l2 * (4.0 * r2 * c * s - 2.0 * r4 * c * s) + 4.0 * l3 * r3 * c2 * s - 2.0 * l4 * r2 * c * s + v * (l * (2.0 * r * s + 8.0 * r3 * c2 * s - 2.0 * r3 * s3) - 10.0 * l2 * r2 * c * s + 2.0 * l3 * r * s);
    a[3] = r2 + l * (2.0 * r3 * c3 - 10.0 * r3 * c * s2) + l2 * (-2.0 * r2 * c2 + 2.0 * r2 * s2 + r4) - 2.0 * l3 * r3 * c + l4 * r2 + v * (l * (-2.0 * r * c - 4.0 * r3 * c3 + 6.0 * r3 * c * s2) + 6.0 * l2 * r2 * c2 - 2.0 * l3 * r * c) + v * v * l2;
    a[2] = -2.0 * r2 * c * s + 8.0 * l * r3 * c2 * s - l2 * (4.0 * r2 * c * s + 2.0 * r4 * c * s) + 4.0 * l3 * r3 * c2 * s - 2.0 * l4 * r2 * c * s + v * (l * (2.0 * r * s - 4.0 * r3 * c2 * s + r3 * s3) - 2.0 * l2 * r2 * c * s + 2.0 * l3 * r * s);
    a[1] = r2 * c2 - 2.0 * l * r3 * c3 + l2 * (2.0 * r2 * c2 + r4 * c2) - 2.0 * l3 * r3 * c3 + l4 * r2 * c2 + v * (-l * (2.0 * r * c + r3 * c * s2) + 2.0 * l2 * r2 * c2 - 2.0 * l3 * r * c) + v * v * l2;

#// solve for roots
    order = 5;

    roots <- polyroot(a)
        tanphi <- Re(roots) 
        rootsimag <- Im(roots)

#/ we should now have 5 values of tan(phi), one for each root!
#/ If there are any imaginary roots, they are NOT REAL SOLUTIONS.
#/ Hence, if inside a caustic, we'll have five real solutions, otherwise 3
        realimage <- abs(rootsimag) < 1e-4
        realcount <- length(rootsimag[realimage == T])

        if (realcount != 3 & realcount != 5) {
            stop("Not 3 or 5 real images!")
#/ by the way. LESS than 3 images are simply dealt with by adding 1 to amplification 
#/ for each missing image
        }

#/ now get r.cos(phi)
    top <- (l2 * r * c - l * r2 * c2 - v * l + r * c) * tanphi * tanphi - r * s * (l2 - 2 * l * r * c + 1.0) * tanphi - l * r2 * s2;
    bottom <- r * (c * tanphi - s) * (r * s + 2 * (l - r * c) * tanphi - r * s * tanphi * tanphi);
    rcosphi <- top/bottom

    #// don't forget to scale back to pvEvent scale!
    scale = sqrt(1.0 + qForScaling);
    imagePositions <- matrix(NA,realcount,2)
    count <- 1
    for (i in 1:5) {
        if (realimage[i]) {
            imagePositions[count,1] <- rcosphi[i]*scale
            imagePositions[count,2] <- rcosphi[i]*tanphi[i]*scale
            count <- count+1
        }
    }

    return(imagePositions)
}


  calcLensEquationDeterminant <- function(a,q, x, y) {
        derx = 0.0;
        dery = 0.0;
        #// lens positions
        px <- c(0.0, a);
        py <- c(0.0, 0.0);
        tq <- c(1.0, q)
        for (i in 1:2) {
            dx = px[i] - x;
            dy = py[i] - y;
            bottom = (dx) * (dx) + (dy) * (dy);
            bottom <- bottom*bottom
            derx <- derx + (dx * dx - dy * dy) / bottom * tq[i];
            dery <- dery +  2.0 * dx * dy * tq[i] / bottom;
        }

        return (1.0 - (derx * derx + dery * dery))
    }


calcAmplificationFromImagePositions <- function(a,q, positions) {
    totalamp = 0.0
    for (i in 1:length(positions)) {
        temp <- calcLensEquationDeterminant(a,q, positions[i,1], positions[i,2])
        totalamp <- totalamp +  1.0 / abs(temp)
    }
    return(totalamp)
}


calcBinaryLensAmplification <- function(day,a,ang,b,q,te,tm,rho=0,psi=0) {
# first obtain source x- and y-position
    pos <- lensSourcePlanePosition(day,a,ang,b,q,te,tm,rho,psi);
# Next get the Asada position parmeterization
    asadaPos <- calcAsadaParams(pos$sx,pos$sy,a,q)
# now calculate the Asada image positions
        positions <- calcImagePositionsUsingAsada(asadaPos, q);
        print(positions)
        #/ calculate amplification from image positions
        amplification <- calcAmplificationFromImagePositions(a,q, positions);
        return(amplification)
}



